import sys
import os
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')))
import yaml
from langgraph.graph import StateGraph
from typing import Dict, Any
from executor.graph.llm_engine import client
# -------------------------
# 1. 示例 YAML 内容
# -------------------------
yaml_sample = """
task: 撰写关于 AI 对初中教学积极影响的综述报告

agents:
  - name: RetrieverAgent
    role: 收集与人工智能在初中教学中的应用、优势、实际案例相关的资料
    outputs: research_notes

  - name: AnalyzerAgent
    role: 对检索到的资料进行分析与分类，提取关键观点与趋势
    inputs: research_notes
    outputs: summarized_findings

  - name: OutlineAgent
    role: 根据分析结果生成报告提纲，划分章节和逻辑结构
    inputs: summarized_findings
    outputs: article_outline

  - name: WritingAgent
    role: 根据提纲撰写文章正文，引用具体数据或案例
    inputs: article_outline, summarized_findings
    outputs: article_draft

  - name: EditingAgent
    role: 优化文章语言、逻辑连贯性，并校对语法错误
    inputs: article_draft
    outputs: final_report

workflow:
  - from: RetrieverAgent
    to: AnalyzerAgent
  - from: AnalyzerAgent
    to: OutlineAgent
  - from: AnalyzerAgent
    to: WritingAgent
  - from: OutlineAgent
    to: WritingAgent
  - from: WritingAgent
    to: EditingAgent
"""


def make_llm_worker(agent_cfg):
    name = agent_cfg["name"]
    role = agent_cfg["role"]
    inputs = agent_cfg.get("inputs", [])
    output = agent_cfg.get("output")
    def worker(state: Dict[str, Any]) -> Dict[str, Any]:
        print(f"[LLM Worker] {name} 正在处理：{role}")

        # 构建 prompt 内容
        context = ""
        for inp in inputs:
            context += f"{inp}: {state['memory'].get(inp, '')}\n"

        messages = [
            {"role": "system", "content": f"你是一个agent，职责如下：{role}"},
            {"role": "user", "content": f"任务背景：{state.get('task', '')}\n相关资料：\n{context}"}
        ]

        # 调用 DeepSeek
        response = client.chat.completions.create(
            model="deepseek-chat",
            messages=messages
        )
        result = response.choices[0].message.content

        # 存入 state
        state["done"].add(name)
        state["results"][output] = result
        return state

    return worker
from typing import TypedDict, Dict

# 完整状态结构
class MyState(TypedDict):
    task: str
    results: dict
# -------------------------
# 4. 构建 LangGraph
# -------------------------
builder = StateGraph(MyState)
def build_graph(yaml_input):
      # -------------------------
  # 2. 加载 YAML 为 Python dict
  # -------------------------
  config = yaml.safe_load(yaml_input)
  agents = {agent['name']: agent for agent in config['agents']}
  workflow = config['workflow']


  # 添加节点（agent）
  for name in agents:
      agent = agents[name]
      node = make_llm_worker(agent)
      builder.add_node(name, node)

  # 添加连接关系（workflow）
  for edge in workflow:
      from_node = edge['from']
      to_node = edge['to']
      builder.add_edge(from_node, to_node)

  # 设置入口和出口
  wf =  config['workflow']
  entry_node = wf[0]['from']
  builder.set_entry_point(entry_node)
  end_node = wf[len(wf) - 1]['to']
  builder.set_finish_point(end_node)  # 假设最后一个

  # 编译为 LangGraph 可运行图
  graph = builder.compile()
  return graph
if __name__ == "__main__":
    build_graph(yaml_sample)